

from sketchc import *

window = Window2((1000, 800))

class Water(Obj):
    def __init__(self, type, size, mesh_size):
        self.type = type
        self.size = size
        self.mesh_size = mesh_size
        self.uv = [(0, 1), (1, 1), (0, 0), (1, 0)]

        self.mode = moderngl.TRIANGLES
        self.t = Turtle()
        self.box = None

        self.prog = {
            "vertex_shader":'''
#version 330 core

in vec3 in_position;
//in vec3 in_normal;
in vec2 in_texcoord_0;

uniform mat4 m_proj;
uniform mat4 m_model;
uniform mat4 m_cam;

out vec4 wPos;     // world
//out vec3 wNormal;     // world
out vec2 uv;

void main() {
    //wNormal = mat3(m_model) * in_normal;
    wPos = m_model * vec4(in_position, 1.0);
    gl_Position = m_proj * m_cam * m_model * vec4(in_position, 1.0);
    uv = in_texcoord_0;
}
            ''',
            "fragment_shader":'''
#version 330 core
out vec4 FragColor;

uniform vec3 lightPos;
uniform vec3 viewPos;
uniform sampler2D texture1;

in vec4 wPos;
//in vec3 wNormal;        // in water, it will be set
in vec2 uv;

float big(float x)
{
    return tan(x * 2 - 1);
}


void main()
{
    float hi = texture2D(texture1, uv).x;
    float x1 = big(texture2D(texture1, uv + vec2(-1, 0) * 0.001).x);
    float x2 = big(texture2D(texture1, uv + vec2(1, 0) * 0.001).x);
    float y1 = big(texture2D(texture1, uv + vec2(0, -1) * 0.001).x);
    float y2 = big(texture2D(texture1, uv + vec2(0, 1) * 0.001).x);
    float dx = (x2 - x1) / 0.002;
    float dy = (y2 - y1) / 0.002;
    vec3 wNormal = cross(vec3(1, dx, 0), vec3(0, dy, -1));

    vec3 objectColor = vec3(0.4, 0.4, 0.4);

    vec3 FragPos = wPos.xyz;
    float ambientStrength = 0.05;
    vec3 ambient = ambientStrength * objectColor;

    vec3 norm = normalize(wNormal);
    vec3 lightDir0 = lightPos - FragPos;
    vec3 lightDir = normalize(lightDir0);
    float diff = max(dot(norm, lightDir) , 0.0);
    //vec3 diffuse = diff * vec3(0.5, 0.5, 0.5) ;
    vec3 diffuse = diff * objectColor ;


    vec3 viewDir = normalize(viewPos - FragPos);
    float spec = 0.0;
    vec3 halfwayDir = normalize(lightDir + viewDir);
    //spec = pow(max(dot(norm, halfwayDir), 0.0), 32.0);
    spec = pow(max(dot(norm, halfwayDir), 0.0), 32.0 * 100);
    //vec3 specular = spec * vec3(0.3);
    float decay = 1 / pow(length(lightDir0), 2);
    vec3 specular = 200 * spec * objectColor * decay;

    //vec3 result = (ambient * 0 + diffuse * 0.1 + specular * 0.9);
    vec3 result = (ambient * 0 + diffuse * 0.5  + specular * 0.5);

    //vec3 red = vec3(1, 0, 0);

    FragColor = vec4(result  , 1.0) *1 + vec4(hi, 0, 0, 1) * 0;

}

            ''',
        }


        self.fbo = None
        self.fbo2 = None
        self.update_pass = None
        self.update_pass2 = None
        self.vao = None

        c = self.c = 0.7
        h = self.h = 1
        k = self.k = 1
        self.kappa = k * c / h
        self.tau = ((c*k) / h) ** 2
        self.alpha = self.tau
        print(self.alpha)


    def __rrshift__(self, obj):
        raise NotImplemented()

    def get_vao(self, ctx):
        c = rect(self.type, *self.size)        # 注意这里的顺序是逆时针序
        p2, p1, p3, p4 = c
        vertices = [p1, p4, p2, p1, p3, p4]
        normals = easy_norm(vertices)
        uv1, uv2, uv3, uv4 = self.uv
        uvs = [uv1, uv4, uv2, uv1, uv3, uv4 ]
        vertices = np.array(vertices, dtype="f4")
        # ~ normals = np.array(normals, dtype="f4")
        uvs = np.array(uvs, dtype="f4")
        a, b, c = vertices[:,0], vertices[:,1], vertices[:,2]
        # ~ d, e, f = normals[:,0], normals[:,1], normals[:,2]
        g, h = uvs[:,0], uvs[:,1]
        # ~ x = np.dstack([a, b, c, d, e, f, g, h])
        x = np.dstack([a, b, c, g, h])
        vbo = ctx.buffer(x.astype("f4"))
        # ~ return ctx.simple_vertex_array(self.prog, vbo,  'in_position', 'in_normal', 'in_texcoord_0')
        return ctx.simple_vertex_array(self.prog, vbo,  'in_position', 'in_texcoord_0')

    def u(self, ctx):
        assert isinstance(ctx, moderngl.Context)
        self.prog = ctx.program(**self.prog)
        self.vao= self.get_vao(ctx)

        self.fbo = EasyFbo(self.mesh_size, ctx, quality=True)
        self.fbo2 = EasyFbo(self.mesh_size, ctx, quality=True)
        # ~ self.fbo.texture.filter = (moderngl.LINEAR_MIPMAP_LINEAR, moderngl.LINEAR)
        self.fbo.filter = moderngl.NEAREST, moderngl.NEAREST
        self.fbo2.filter = moderngl.NEAREST, moderngl.NEAREST
        update_prog = {
            "vertex_shader":'''
#version 330 core

in vec3 in_position;
in vec2 in_texcoord_0;

uniform mat4 all_mat;

out vec2 uv;

void main() {
    gl_Position = all_mat * vec4(in_position, 1.0);
    uv = vec2(in_texcoord_0.x, 1 - in_texcoord_0.y);
}
            ''',
            "fragment_shader":'''
#version 330 core
out vec4 fragColor;
uniform sampler2D input_texture;
uniform float alpha;
uniform vec2 mesh_size;
in vec2 uv;

float small(float x)
{
    return (atan(x) + 1) / 2;
}

float big(float x)
{
    return tan(x * 2 - 1);
}

float getx(vec2 uv)
{
    return big(texture2D(input_texture, uv).x);
}

float gety(vec2 uv)
{
    return big(texture2D(input_texture, uv).y);
}

void main()
{
    vec2 d = vec2(1 / mesh_size.x, 1 / mesh_size.y);
    vec4 r = vec4(1);
    //r.z = gety(uv);
    r.y = getx(uv);
    r.x = getx(vec2(uv.x - d.x, uv.y)) +
                getx(vec2(uv.x + d.x, uv.y)) +
                getx(vec2(uv.x, uv.y - d.y)) +
                getx(vec2(uv.x, uv.y + d.y)) -
                getx(uv) * 4;
    r.x = r.x * alpha +
                      getx(uv) * 2 -
                      gety(uv);

    //fragColor = r * 0 + vec4(sin(uv.x * 100) * sin(uv.y * 100 ), 0.2, 0.2, 1);
    fragColor = vec4(small(r.x), small(r.y), small(r.z), 0) ;
}

            ''',
        }
        self.update_pass = Pass(self.mesh_size, ctx, ctx.program(**update_prog), self.fbo)
        self.update_pass.prog['mesh_size'].write(np.array(self.mesh_size, dtype="f4"))
        self.update_pass.prog['alpha'].value = self.alpha

    def update(self):
        self.fbo2.use()
        self.fbo2.ctx.clear(0.5, 0.5, 0.5)
        self.update_pass.input_fbo = self.fbo
        self.update_pass.render()       # 此时 fbo1.texture.use
        self.fbo, self.fbo2 = self.fbo2, self.fbo
        # ~ self.point = 1 - self.point
        # ~ print("hi")


    def draw(self, camera):
        if self.vao is None:
            print("没准备好绘制", self)
            return

        self.prog['m_proj'].write(camera.p())
        self.prog['m_model'].write(self.m())
        self.prog['m_cam'].write(camera.v())

        self.prog['lightPos'].write(U3(0, 0, 1))
        self.prog['viewPos'].write(camera.pos)

        self.prog['texture1'].value = 1
        self.fbo.texture.use(1)
        # ~ print(self.fbo.texture)

        self.vao.render(self.mode)

    def draw_normal(self, camera):
        return



mesh_size =  (100, 100)
mesh_size =  (1000, 1000)
water = Water("xy", (10, 10), mesh_size)


water.u(window.ctx)
water >> window


def small(x):
    return (np.arctan(x) + 1) / 2



camera = Camera(U3(0, -5, 1))
camera >> window        # 识别为camera则自动作为主camera

walker = Walker(camera)
walker >> window

@id(lambda f: window.update_callback.append(f))
def f(window, dt):
    water.update()


tempvar = [0]
@id(lambda f: window.mouse_down_callback.append(f))
def f(window, e):
    # ~ data = water.main_texture.read()
    data = water.fbo.texture.read()
    mat = np.frombuffer(data, dtype="f4").reshape((mesh_size[0], mesh_size[1], 4)).copy()
    print(type(mat))
    p = np.random.random(2) * 2 - 1
    p = p * 200 + np.array([500, 500])
    p = p.astype(np.int32)
    r = np.random.random(mesh_size) * 2 - 1
    r = r / 1000
    r = small(r)
    # ~ mat[:, :, 3] = r
    x = tempvar[0]
    mat[p[0] - 10: p[0] + 10, p[1] - 10: p[1] + 10, 0] = 0.0002
    water.fbo.texture.write(mat)



clock = pygame.time.Clock()
while True:
    dt = clock.tick(30)
    window.update(dt)
    for e in pygame.event.get():
        if e.type == pygame.QUIT:
            exit()
        window.dispatch(e)
    window.render()   # screen 操作需要写在window的槽里
    pygame.display.flip()

